#ifndef BTL_ROBUSTESTIMATION_PREEMPTIVERANSAC
#define BTL_ROBUSTESTIMATION_PREEMPTIVERANSAC

#include "RANSAC.hpp"
#include <btl/Maths/RandomSampler.hpp>

#include <boost/optional.hpp>
#include <boost/random/random_number_generator.hpp>
#include <boost/random/mersenne_twister.hpp>

#include <vector>
#include <algorithm>
#include <cassert>

namespace btl {
namespace robustestimation {

template <class ModelType, int SampleSize, class RNG = boost::random_number_generator<boost::mt19937> >
class PreemptiveRANSAC : public RobustModelFinder<ModelType, RNG>
{
   public:
      PreemptiveRANSAC(int numSamples, double inlierThreshold, int scoreBlockSize);

      virtual boost::optional<ModelType> findModel(const ModelSearchSpace<ModelType>& searchSpace, RNG& rng);

   private:
      struct Score
      {
         explicit Score(int model): model(model), cost(0.0) {}
         Score(int model, double cost): model(model), cost(cost) {}

         int model;
         double cost;

         friend bool operator<(const Score& a, const Score& b)
         { return a.cost < b.cost; }
      };

      double _inlierThreshold;
      int _numSamples;
      int _scoreBlockSize;
};

} //robustestimation
} //btl


namespace btl
{

using robustestimation::PreemptiveRANSAC;

} //btl

// ====================================================================
// === Implementation

namespace btl {
namespace robustestimation {

template <class ModelType, int SampleSize, class RNG>
inline PreemptiveRANSAC<ModelType,SampleSize,RNG>::PreemptiveRANSAC(
      int numSamples, double inlierThreshold, int scoreBlockSize):
   _inlierThreshold(inlierThreshold), _numSamples(numSamples), _scoreBlockSize(scoreBlockSize)
{
   // FIXME: block size should be calculated so that at f(N) or so, we have one hypothesis left
}

template <class ModelType, int SampleSize, class RNG>
inline boost::optional<ModelType> PreemptiveRANSAC<ModelType,SampleSize,RNG>::findModel(
      const ModelSearchSpace<ModelType>& searchSpace, RNG& rng)
{
   const int N = searchSpace.numObservations();

   std::vector<int> obsIndexes;
   obsIndexes.reserve(N);
   for (int i = 0; i < N; ++i)
      obsIndexes.push_back(i);

   std::vector<ModelType> models;
   models.reserve(_numSamples);

   // generate models
   for (int i = 0; i < _numSamples; ++i)
   {
      // select a sample
      btl::shuffleMofN(
         obsIndexes.begin(),
         obsIndexes.begin() + SampleSize,
         obsIndexes.end(),
         rng);

      // add the model from that sample
      searchSpace.buildModels(&obsIndexes[0], SampleSize, models);
   }

   //std::cout << "Have " << N << " data points to use\n";
   //std::cout << "Generated " << models.size() << " models\n";

   std::vector<Score> scores;
   scores.reserve(models.size());
   for (int i = 0; i < int(models.size()); ++i)
      scores.push_back(Score(i));

   // shuffle all observations to avoid bias
   std::random_shuffle(obsIndexes.begin(), obsIndexes.end(), rng);

   const int M = models.size();
   int keep = M;
   for (int i = 0; (keep > 1) && (i < obsIndexes.size()); ++i) {
      int newKeep = M >> (i / _scoreBlockSize);

      // partition to put the best scores at the front
      if (newKeep != keep)
      {
         //std::cout << "At step " << i << " dropping to " << newKeep << " models\n";
         assert(newKeep < keep);
         std::nth_element(scores.begin(), scores.begin() + newKeep, scores.begin() + keep);
         keep = newKeep;

         //std::cout << "First and last score within the new range:\n";
         //int j = 0;
         //std::cout << "[" << j << "]; model " << scores[j].model << "; cost = " << scores[j].cost << "\n";
         //j = keep-1;
         //std::cout << "[" << j << "]; model " << scores[j].model << "; cost = " << scores[j].cost << "\n";
      }

      // score all remaining hypotheses on this observation
      const int obsIdx = obsIndexes[i];
      for (int j = 0; j < keep; ++j)
      {
         double c = searchSpace.errorMetric(models[scores[j].model], obsIdx);
         assert(c >= 0.0);
         c = std::min(c, _inlierThreshold); // all outliers cost the same

         scores[j].cost += c;
      }
   }

   return models[scores[0].model];
}

} //robustestimation
} //btl

#endif //BTL_ROBUSTESTIMATION_PREEMPTIVERANSAC
